import numpy as np
import matplotlib.pyplot as plt
from typing import List, Tuple
from fg_grc_plate import (GrapheneReinforcedComposite, MaterialProperties, 
                         GrapheneGeometry, DistributionPattern)
from laminated_plate import LaminatedPlate, Layer
from vibration_analysis import PlateVibrationSolver, BoundaryCondition

class ValidationStudy:
    """
    Validation studies reproducing results from the paper.
    """
    
    def __init__(self):
        # Material properties from Table 1
        self.epoxy = MaterialProperties(E=3.0, nu=0.34, rho=1200)
        self.graphene = MaterialProperties(E=1010, nu=0.186, rho=1060)
        self.gpl_geometry = GrapheneGeometry(
            length=2.5e-6,    # 2.5 μm
            width=1.5e-6,     # 1.5 μm
            thickness=1.5e-9  # 1.5 nm
        )
        
        # Create GRC material
        self.grc = GrapheneReinforcedComposite(self.epoxy, self.graphene, self.gpl_geometry)
    
    def isotropic_plate_validation(self) -> Tuple[List[float], List[float]]:
        """
        Validate against analytical solution for isotropic plate.
        Reproduces Table 2 from the paper.
        """
        print("Validation Study 1: Isotropic Square Plate")
        print("=" * 50)
        
        # Test different thickness ratios
        h_over_a_ratios = [0.01, 0.1, 0.2]
        plate_size = 0.1  # 10 cm
        
        numerical_results = []
        analytical_results = []
        
        for h_over_a in h_over_a_ratios:
            thickness = h_over_a * plate_size
            
            # Create isotropic plate
            layers = [Layer(thickness=thickness, material=self.epoxy)]
            plate = LaminatedPlate(layers, length=plate_size, width=plate_size)
            
            # Solve numerically
            solver = PlateVibrationSolver(plate, n_points_x=11, n_points_y=11)
            bc_ss = (BoundaryCondition.SIMPLY_SUPPORTED, BoundaryCondition.SIMPLY_SUPPORTED)
            freqs, _ = solver.solve_eigenvalue_problem(bc_ss, bc_ss, n_modes=1)
            
            # Calculate dimensionless frequency
            dim_freq_numerical = solver.dimensionless_frequency(freqs[0], self.epoxy)
            
            # Analytical solution for SSSS square plate (first mode)
            # ω = π²√(D/(ρh)) * √(2)/a²
            D = self.epoxy.E * thickness**3 / (12 * (1 - self.epoxy.nu**2))
            omega_analytical = np.pi**2 * np.sqrt(2 * D / (self.epoxy.rho * thickness)) / plate_size**2
            dim_freq_analytical = omega_analytical * plate_size**2 / thickness * \
                                np.sqrt(self.epoxy.rho / self.epoxy.E)
            
            numerical_results.append(dim_freq_numerical)
            analytical_results.append(dim_freq_analytical)
            
            error = abs(dim_freq_numerical - dim_freq_analytical) / dim_freq_analytical * 100
            
            print(f"h/a = {h_over_a:.2f}:")
            print(f"  Numerical:  Ω = {dim_freq_numerical:.4f}")
            print(f"  Analytical: Ω = {dim_freq_analytical:.4f}")
            print(f"  Error:      {error:.2f}%")
            print()
        
        return numerical_results, analytical_results
    
    def gpl_weight_fraction_study(self) -> Tuple[List[float], List[float]]:
        """
        Study effect of GPL weight fraction.
        Reproduces Figure 2 from the paper.
        """
        print("\nValidation Study 2: Effect of GPL Weight Fraction")
        print("=" * 50)
        
        # Convert weight fraction to volume fraction (approximate)
        weight_fractions = np.linspace(0, 0.05, 11)  # 0 to 5%
        volume_fractions = []
        frequencies = []
        
        plate_size = 0.1  # 10 cm
        total_thickness = 2.0e-3  # 2 mm
        
        for W_G in weight_fractions:
            # Approximate conversion from weight to volume fraction
            # W_G ≈ V_G * ρ_G / ρ_c (simplified)
            V_G = W_G * self.epoxy.rho / self.graphene.rho
            volume_fractions.append(V_G)
            
            # Create three-layer [0°/GRC/90°] plate
            layer_thickness = total_thickness / 3
            layers = [
                Layer(thickness=layer_thickness, material=self.epoxy, angle=0.0),
                Layer(thickness=layer_thickness, material=self.epoxy, is_grc=True,
                      grc_material=self.grc, grc_pattern=DistributionPattern.UD,
                      grc_volume_fraction=V_G, grc_exponent=1.0),
                Layer(thickness=layer_thickness, material=self.epoxy, angle=90.0)
            ]
            
            plate = LaminatedPlate(layers, length=plate_size, width=plate_size)
            solver = PlateVibrationSolver(plate, n_points_x=9, n_points_y=9)
            
            bc_ss = (BoundaryCondition.SIMPLY_SUPPORTED, BoundaryCondition.SIMPLY_SUPPORTED)
            freqs, _ = solver.solve_eigenvalue_problem(bc_ss, bc_ss, n_modes=1)
            
            frequencies.append(freqs[0])
            
            print(f"W_G = {W_G*100:.1f}%, V_G = {V_G*100:.2f}%, ω = {freqs[0]:.2f} rad/s")
        
        return weight_fractions, frequencies
    
    def distribution_pattern_study(self) -> dict:
        """
        Study effect of graphene distribution patterns.
        Reproduces Table 3 from the paper.
        """
        print("\nValidation Study 3: Effect of Distribution Patterns")
        print("=" * 50)
        
        patterns = {
            'UD': DistributionPattern.UD,
            'FG-V': DistributionPattern.FG_V,
            'FG-X': DistributionPattern.FG_X,
            'FG-O': DistributionPattern.FG_O
        }
        
        results = {}
        
        # Plate parameters
        plate_size = 0.1  # 10 cm
        h_over_a = 0.1
        thickness = h_over_a * plate_size
        W_G = 0.01  # 1% weight fraction
        V_G = W_G * self.epoxy.rho / self.graphene.rho  # Convert to volume fraction
        
        for pattern_name, pattern in patterns.items():
            # Create single GRC layer
            layers = [Layer(thickness=thickness, material=self.epoxy, is_grc=True,
                          grc_material=self.grc, grc_pattern=pattern,
                          grc_volume_fraction=V_G, grc_exponent=1.0)]
            
            plate = LaminatedPlate(layers, length=plate_size, width=plate_size)
            solver = PlateVibrationSolver(plate, n_points_x=9, n_points_y=9)
            
            bc_ss = (BoundaryCondition.SIMPLY_SUPPORTED, BoundaryCondition.SIMPLY_SUPPORTED)
            freqs, _ = solver.solve_eigenvalue_problem(bc_ss, bc_ss, n_modes=5)
            
            # Calculate dimensionless frequencies
            dim_freqs = [solver.dimensionless_frequency(freq, self.epoxy) for freq in freqs]
            results[pattern_name] = dim_freqs
            
            print(f"{pattern_name}:")
            for i, dim_freq in enumerate(dim_freqs):
                print(f"  Mode {i+1}: {dim_freq:.3f}")
            print()
        
        return results
    
    def boundary_condition_study(self) -> dict:
        """
        Study effect of boundary conditions.
        """
        print("\nValidation Study 4: Effect of Boundary Conditions")
        print("=" * 50)
        
        boundary_conditions = {
            'SSSS': ((BoundaryCondition.SIMPLY_SUPPORTED, BoundaryCondition.SIMPLY_SUPPORTED),
                     (BoundaryCondition.SIMPLY_SUPPORTED, BoundaryCondition.SIMPLY_SUPPORTED)),
            'CCCC': ((BoundaryCondition.CLAMPED, BoundaryCondition.CLAMPED),
                     (BoundaryCondition.CLAMPED, BoundaryCondition.CLAMPED)),
        }
        
        results = {}
        
        # Create reference plate
        plate_size = 0.1
        thickness = 2.0e-3
        layers = [Layer(thickness=thickness, material=self.epoxy)]
        plate = LaminatedPlate(layers, length=plate_size, width=plate_size)
        solver = PlateVibrationSolver(plate, n_points_x=9, n_points_y=9)
        
        for bc_name, (bc_x, bc_y) in boundary_conditions.items():
            freqs, _ = solver.solve_eigenvalue_problem(bc_x, bc_y, n_modes=3)
            dim_freqs = [solver.dimensionless_frequency(freq, self.epoxy) for freq in freqs]
            results[bc_name] = dim_freqs
            
            print(f"{bc_name}:")
            for i, dim_freq in enumerate(dim_freqs):
                print(f"  Mode {i+1}: {dim_freq:.3f}")
            print()
        
        return results

def create_validation_plots(validation: ValidationStudy):
    """Create plots for validation studies."""
    
    # Study 1: Validation plot
    numerical, analytical = validation.isotropic_plate_validation()
    
    # Study 2: GPL weight fraction effect
    weight_fractions, frequencies = validation.gpl_weight_fraction_study()
    
    # Study 3: Distribution patterns
    pattern_results = validation.distribution_pattern_study()
    
    # Study 4: Boundary conditions
    bc_results = validation.boundary_condition_study()
    
    # Create plots
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))
    
    # Plot 1: Validation
    h_over_a = [0.01, 0.1, 0.2]
    ax1.plot(h_over_a, numerical, 'bo-', label='Numerical (DQ)')
    ax1.plot(h_over_a, analytical, 'rs--', label='Analytical')
    ax1.set_xlabel('h/a')
    ax1.set_ylabel('Dimensionless Frequency Ω')
    ax1.set_title('Validation: Isotropic Square Plate (SSSS)')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # Plot 2: GPL weight fraction effect
    freq_enhancement = np.array(frequencies) / frequencies[0]
    ax2.plot(np.array(weight_fractions)*100, freq_enhancement, 'g-', linewidth=2)
    ax2.set_xlabel('GPL Weight Fraction (%)')
    ax2.set_ylabel('Frequency Enhancement Factor')
    ax2.set_title('Effect of GPL Weight Fraction')
    ax2.grid(True, alpha=0.3)
    
    # Plot 3: Distribution patterns (first mode)
    patterns = list(pattern_results.keys())
    first_mode_freqs = [pattern_results[p][0] for p in patterns]
    ax3.bar(patterns, first_mode_freqs, color=['blue', 'red', 'green', 'orange'])
    ax3.set_ylabel('Dimensionless Frequency')
    ax3.set_title('Effect of GPL Distribution Patterns (Mode 1)')
    ax3.grid(True, alpha=0.3, axis='y')
    
    # Plot 4: Boundary conditions
    bc_names = list(bc_results.keys())
    mode1_freqs = [bc_results[bc][0] for bc in bc_names]
    ax4.bar(bc_names, mode1_freqs, color=['cyan', 'magenta'])
    ax4.set_ylabel('Dimensionless Frequency')
    ax4.set_title('Effect of Boundary Conditions (Mode 1)')
    ax4.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    plt.savefig('/home/linuxuser/code/work/NS/validation_results.png', dpi=300, bbox_inches='tight')
    plt.show()

if __name__ == "__main__":
    # Run all validation studies
    validation = ValidationStudy()
    
    print("REPRODUCING RESULTS FROM:")
    print("Dynamic Analysis of Composite Laminated Plates with")
    print("a Functionally Graded Graphene-Reinforced Layer")
    print("(Wang et al., 2015)")
    print("=" * 60)
    
    # Run validation studies
    validation.isotropic_plate_validation()
    validation.gpl_weight_fraction_study()
    validation.distribution_pattern_study()
    validation.boundary_condition_study()
    
    # Create plots
    create_validation_plots(validation)
    
    print("\nValidation studies completed successfully!")
    print("Results saved to validation_results.png")